"""
A simple example demonstrating the usage of Sonetto framework
"""
import os
from sonetto.agents.user import UserProxy
from sonetto.agents.chat import ChatAgent
from sonetto.agents.tool import ToolAgent
from sonetto.core.scheduler import Scheduler
from sonetto.core.toolset import McpToolset
from sonetto.core.message import Message, MessageType

def get_weather(city: str) -> str:
    """模拟天气查询工具"""
    return f"{city}的天气是晴朗，温度25℃"

def main():
    # 获取API配置
    api_key = os.getenv("API_KEY")
    if not api_key:
        raise ValueError("请设置API_KEY环境变量")
    
    # 可选：设置自定义base_url（例如使用智谱AI）
    base_url = os.getenv("API_BASE_URL")

    # 创建调度器
    scheduler = Scheduler()

    # 创建用户代理
    user = UserProxy()
    scheduler.register_agent(user)

    # 创建聊天代理
    chat_agent = ChatAgent(
        agent_id="chat_assistant",
        api_key=api_key,
        system_prompt="你是一个友好的助手，可以跟用户聊天。",
        base_url=base_url  # 如果设置了自定义base_url，则使用它
    )
    scheduler.register_agent(chat_agent)

    # 创建工具集
    toolset = McpToolset()

    # 注册工具
    @toolset.register_tool(
        name="get_weather",
        description="查询指定城市的天气情况"
    )
    def weather_tool(city: str) -> str:
        return get_weather(city)

    # 创建工具代理
    tool_agent = ToolAgent(
        agent_id="tool_assistant",
        api_key=api_key,
        toolset=toolset,
        base_url=base_url  # 如果设置了自定义base_url，则使用它
    )
    scheduler.register_agent(tool_agent)

    print("=== Sonetto 多Agent聊天示例 ===")
    print(f"当前使用的API地址: {base_url or 'https://api.openai.com/v1'}")
    print("可用命令：")
    print("- 输入 'chat:消息' 与聊天助手对话")
    print("- 输入 'tool:消息' 与工具助手对话")
    print("- 输入 'exit' 退出")
    print("=============================")

    while True:
        try:
            user_input = input("\n你: ").strip()
            if user_input.lower() == 'exit':
                break

            if user_input.startswith("chat:"):
                # 发送消息给聊天助手
                content = user_input[5:].strip()
                message = Message(
                    type=MessageType.TEXT,
                    content=content,
                    sender=user.agent_id,
                    receiver=chat_agent.agent_id
                )
            elif user_input.startswith("tool:"):
                # 发送消息给工具助手
                content = user_input[5:].strip()
                message = Message(
                    type=MessageType.TEXT,
                    content=content,
                    sender=user.agent_id,
                    receiver=tool_agent.agent_id
                )
            else:
                print("请使用 'chat:' 或 'tool:' 前缀指定接收者")
                continue

            # 发送消息并获取响应
            response = scheduler.step(message)
            if response:
                scheduler.step(response)

        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main()
